package cn.iocoder.yudao.module.dify.service.BatchUploadKnowledgeBase;

import cn.hutool.core.io.FileUtil;
import cn.iocoder.yudao.module.dify.config.DifyKnowledgeBaseManager;
import cn.iocoder.yudao.module.dify.util.FileStructureUtil;
import io.github.imfangs.dify.client.DifyDatasetsClient;
import io.github.imfangs.dify.client.exception.DifyApiException;
import io.github.imfangs.dify.client.model.datasets.*;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections4.ListUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.io.File;
import java.io.IOException;
import java.util.*;
import java.util.stream.Collectors;

@Slf4j
@Service
public class BatchUploadKnowledgeBaseServiceImpl implements BatchUploadKnowledgeBaseService {

    @Autowired
    private DifyKnowledgeBaseManager difyKnowledgeBaseManager;

    @Autowired
    private PDFProcessService pdfProcessService;

    private String datasetKey = "batchUpload";

    @Override
    public Boolean folderBatchUpload() throws DifyApiException, IOException {
        batchUploadKnowledgeBase(null);
        return true;
    }

    /**
     * 根据指定路径遍历文件夹，并上传到指定知识库
     */
    public void batchUploadKnowledgeBase(String path) throws IOException, DifyApiException {
        path = FileUtil.normalize("D:\\puyuanData\\ISOIEC");
        // 第一步：收集所有文件夹并创建知识库
        Map<String, String> folderToDatasetIdMap = new HashMap<>();
        // 获取文件夹及文件映射关系
        // 限制文件类型
        List<String> extensions = Arrays.asList("txt", "pdf", "docx", "doc");
        Map<String, List<File>> folderFilesMap = FileStructureUtil.getFolderFilesMap(path, extensions);
        // 获取所有知识库
        List<DatasetListResponse.DatasetInfo> allDatasets = getAllDatasets();
        if (allDatasets != null) {
            for (DatasetListResponse.DatasetInfo dataset : allDatasets) {
                folderToDatasetIdMap.put(dataset.getName(), dataset.getId());
            }
        }
        // 先创建所有文件夹对应的知识库
        for (String folderName : folderFilesMap.keySet()) {
            // 避免重复创建同名知识库
            if (!folderToDatasetIdMap.containsKey(folderName)) {
                // 创建知识库
                String datasetId = createKnowledgeBaseAndReturnId(folderName);
                folderToDatasetIdMap.put(folderName, datasetId);
            }
        }
        // 第二步：上传文件到对应的知识库
        // 比较文件是否已存在知识库
        for (String folderName : folderFilesMap.keySet()) {
            String datasetId = folderToDatasetIdMap.get(folderName);
            List<File> filePathList = folderFilesMap.get(folderName);
            Map<String, File> fileNameToFileMap = filePathList.stream()
                    .collect(Collectors.toMap(File::getName, file -> file));
            List<String> newFolderFileName = filePathList.stream().map(File::getName).collect(Collectors.toList());
            List<DocumentListResponse.DocumentInfo> documentListOfTheKnowledgeBase = getDocumentListOfTheKnowledgeBase(datasetId, null);
            //documentListOfTheKnowledgeBase.name转map
            Map<String, String> documentNameMap = documentListOfTheKnowledgeBase.stream()
                    .collect(Collectors.toMap(DocumentListResponse.DocumentInfo::getName, DocumentListResponse.DocumentInfo::getId));
            //知识库中获取的已存在文档名称列表
            List<String> aListOfDocumentNamesAlreadyExists = documentListOfTheKnowledgeBase.stream().map(DocumentListResponse.DocumentInfo::getName).collect(Collectors.toList());
            //需要上传到知识库的新文件列表
            List<String> subtract = ListUtils.subtract(newFolderFileName, aListOfDocumentNamesAlreadyExists);
            for (String fileName : subtract) {
                File file = fileNameToFileMap.get(fileName);
                if (file != null) {
                    if (documentNameMap.containsKey(FileStructureUtil.pdfFileWithTheSameNameAsTxtFile(file.getName()))) {
                        //已在dify有同名的txt文件,跳过
                        continue;
                    }
                    if (pdfProcessService.needOCR(file)) {
                        // 对需要OCR的PDF文件进行异步处理
                        log.info("对需要OCR的PDF文件进行异步处理: " + file.getName());
                        pdfProcessService.addFileToOCRQueue(file, datasetId);
                        continue;
                    }
                    if (documentNameMap.containsKey(FileStructureUtil.docFileWithTheSameNameAsDocxFile(file.getName()))) {
                        //已在dify有同名的docx文件,跳过
                        continue;
                    }
                    if (FileStructureUtil.isDocFile(file.getName())) {
                        log.info("对doc文件进行异步处理: " + file.getName());
                        pdfProcessService.addFileToDocxQueue(file, datasetId);
                        continue;
                    }
                    log.info("直接上传文件到知识库: " + file.getName());
                    uploadFileToKnowledgeBase(file, datasetId);
                }
            }
        }

    }

    /**
     * 创建知识库并返回ID
     */
    public String createKnowledgeBaseAndReturnId(String folderName) {
        DifyDatasetsClient datasetsClient = difyKnowledgeBaseManager.getClient(datasetKey);
        try {
            CreateDatasetRequest createRequest = CreateDatasetRequest.builder()
                    .name(folderName)
                    .description("这是一个批量上传" + folderName + "知识库")
                    .indexingTechnique("high_quality")
                    .permission("all_team_members")
                    .provider("vendor")
                    .build();

            DatasetResponse response = datasetsClient.createDataset(createRequest);
            System.out.println("创建知识库成功: " + folderName + "，ID: " + response.getId());
            return response.getId();
        } catch (Exception e) {
            System.err.println("创建知识库失败: " + e.getMessage());
            e.printStackTrace();
            return null;
        }
    }

    /**
     * 获取知识库详情
     *
     * @param testDatasetId
     * @throws IOException
     * @throws DifyApiException
     */
    public void getDataset(String testDatasetId) throws IOException, DifyApiException {
        // 跳过测试如果没有测试知识库
        if (testDatasetId == null) {
            System.out.println("跳过测试，因为没有测试知识库");
            return;
        }
        DifyDatasetsClient datasetsClient = difyKnowledgeBaseManager.getClient(datasetKey);
        // 获取知识库详情
        DatasetResponse response = datasetsClient.getDataset(testDatasetId);
        // 打印知识库详情
        System.out.println("知识库详情:");
        System.out.println("  ID: " + response.getId());
        System.out.println("  名称: " + response.getName());
        System.out.println("  描述: " + response.getDescription());
        System.out.println("  权限: " + response.getPermission());
        System.out.println("  索引技术: " + response.getIndexingTechnique());
        System.out.println("  文档数量: " + response.getDocumentCount());
        System.out.println("  字数: " + response.getWordCount());
        System.out.println("  嵌入模型: " + response.getEmbeddingModel());
        System.out.println("  嵌入模型提供商: " + response.getEmbeddingModelProvider());

        if (response.getRetrievalModelDict() != null) {
            System.out.println("  检索模型:");
            System.out.println("    搜索方法: " + response.getRetrievalModelDict().getSearchMethod());
            System.out.println("    Top-K: " + response.getRetrievalModelDict().getTopK());
            System.out.println("    重排序启用: " + response.getRetrievalModelDict().getRerankingEnable());
        }
    }

    private ProcessRule.Rules createProcessRules() {
        ProcessRule.Rules rules = new ProcessRule.Rules();
        // 分段规则
        ProcessRule.Segmentation segmentation = new ProcessRule.Segmentation("\n\n", 200);

        // 预处理规则
        ProcessRule.PreProcessingRule removeExtraSpaces = ProcessRule.PreProcessingRule.builder()
                .id("remove_extra_spaces")
                .enabled(true).build();
        List<ProcessRule.PreProcessingRule> preProcessingRules = new ArrayList<>();
        preProcessingRules.add(removeExtraSpaces);
        //子分段规则
        ProcessRule.SubchunkSegmentation subchunkSegmentation = ProcessRule.SubchunkSegmentation.builder()
                .maxTokens(100)
                .chunkOverlap(40)
                .separator("\n")
                .build();
        rules.setPreProcessingRules(preProcessingRules);
        rules.setParentMode("paragraph");
        rules.setSegmentation(segmentation);
        rules.setSubchunkSegmentation(subchunkSegmentation);

        return rules;
    }


    /**
     * 文件上传到知识库
     *
     * @param file
     * @param datasetId
     * @throws IOException
     * @throws DifyApiException
     */
    @Override
    public void uploadFileToKnowledgeBase(File file, String datasetId) throws IOException, DifyApiException {
        DifyDatasetsClient datasetsClient = difyKnowledgeBaseManager.getClient(datasetKey);
        // 跳过测试如果没有测试知识库
        if (datasetId == null) {
            System.out.println("跳过测试，因为没有测试知识库");
            return;
        }
        if (!file.exists()) {
            System.out.println("文件不存在，跳过测试");
            return;
        }

        RetrievalModel retrievalModel = new RetrievalModel();
        retrievalModel.setSearchMethod("hybrid_search");
        retrievalModel.setRerankingEnable(true);
        retrievalModel.setTopK(2);
        retrievalModel.setScoreThreshold(0.7f);
        retrievalModel.setScoreThresholdEnabled(true);
        ProcessRule.Rules processRules = createProcessRules();

        // 创建文档请求
        CreateDocumentByFileRequest request = CreateDocumentByFileRequest.builder()
                .indexingTechnique("high_quality")
                .docForm("text_model")
                .docLanguage("Chinese")
                .indexingTechnique("high_quality")
                .retrievalModel(retrievalModel)
                // 没有这里的设置，会500报错，服务器内部错误
                .processRule(ProcessRule.builder().mode("custom")
                        .rules(processRules)
                        .build())
                .build();

        // 发送请求
        DocumentResponse response = datasetsClient.createDocumentByFile(datasetId, request, file);
        // 保存文档ID用于后续测试
        datasetId = response.getDocument().getId();
    }

    /**
     * 获取知识库列表
     * 处理分页获取所有数据集
     */
    public List<DatasetListResponse.DatasetInfo> getAllDatasets() throws DifyApiException, IOException {
        List<DatasetListResponse.DatasetInfo> allDatasets = new ArrayList<>();
        DifyDatasetsClient datasetsClient = difyKnowledgeBaseManager.getClient(datasetKey);
        int page = 1;
        int limit = 100;
        boolean hasMore = true;

        while (hasMore) {
            DatasetListResponse response = datasetsClient.getDatasets(page, limit);
            // 将当前页的数据添加到总列表中
            allDatasets.addAll(response.getData());
            // 更新分页状态
            hasMore = response.getHasMore();
            page++;
        }
        return allDatasets;
    }

    /**
     * 获取知识库的文档列表
     * keyword
     * 在文档名称中搜索的关键词。
     */
    public List<DocumentListResponse.DocumentInfo> getDocumentListOfTheKnowledgeBase(String datasetId, String keyword) throws IOException, DifyApiException {
        List<DocumentListResponse.DocumentInfo> documentInfos = new ArrayList<>();
        DifyDatasetsClient datasetsClient = difyKnowledgeBaseManager.getClient(datasetKey);
        int page = 1;
        int limit = 100;
        boolean hasMore = true;
        while (hasMore) {
            DocumentListResponse response = datasetsClient.getDocuments(datasetId, keyword, page, limit);
            // 处理当前页的数据
            // 将当前页的数据添加到总列表中
            documentInfos.addAll(response.getData());
            // 更新分页状态
            hasMore = response.getHasMore();
            page++;
        }
        return documentInfos;
    }

    /**
     * 从文本创建文档
     *
     * @param datasetId
     * @param text
     * @param name
     * @param docType
     * @return
     * @throws DifyApiException
     * @throws IOException
     */
    public Boolean documentCreateByText(String datasetId, String text, String name, String docType) throws DifyApiException, IOException {
        DifyDatasetsClient datasetsClient = difyKnowledgeBaseManager.getClient(datasetKey);
        RetrievalModel retrievalModel = new RetrievalModel();
        retrievalModel.setSearchMethod("hybrid_search");
        retrievalModel.setRerankingEnable(true);
        retrievalModel.setTopK(2);
        retrievalModel.setScoreThreshold(0.7f);
        retrievalModel.setScoreThresholdEnabled(true);
        ProcessRule.Rules processRules = createProcessRules();

        // 创建文档请求
        CreateDocumentByTextRequest request = CreateDocumentByTextRequest.builder()
                .name(name)
                .text(text)
                .docType(docType)
                .indexingTechnique("high_quality")
                .docForm("text_model")
                .docLanguage("Chinese")
                .indexingTechnique("high_quality")
                .retrievalModel(retrievalModel)
                // 没有这里的设置，会500报错，服务器内部错误
                .processRule(ProcessRule.builder().mode("custom")
                        .rules(processRules)
                        .build())
                .build();
        DocumentResponse datasetId1 = datasetsClient.createDocumentByText(datasetId, request);
        return true;
    }

    /**
     * 处理OCR并上传到知识库
     *
     * @param file
     * @param datasetId
     */
    public void processFileWithOCR(File file, String datasetId) throws Exception {
        // 对需要OCR的PDF文件进行
        // 执行OCR任务，返回生成的txt文件
        File txtFile = pdfProcessService.performOCR(file);
        if (txtFile != null && txtFile.exists()) {
            uploadFileToKnowledgeBase(txtFile, datasetId);
            System.out.println("OCR结果文件已上传到知识库: " + txtFile.getName());
        }
    }


}
